{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
    "os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = 'prepare/mesolitica-tpu.json'\n",
    "b2_application_key_id = os.environ['b2_application_key_id']\n",
    "b2_application_key = os.environ['b2_application_key']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "client = storage.Client()\n",
    "bucket = client.bucket('mesolitica-tpu-general')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "best = '1020200'\n",
    "directory = 't5-tiny-paraphrase'\n",
    "!rm -rf output out {directory}\n",
    "!mkdir {directory}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = best\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00000-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.data-00001-of-00002')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.index')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.index')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/model.ckpt-{model}.meta')\n",
    "blob.download_to_filename(f'{directory}/model.ckpt-{model}.meta')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/checkpoint')\n",
    "blob.download_to_filename(f'{directory}/checkpoint')\n",
    "\n",
    "blob = bucket.blob(f'{directory}/operative_config.gin')\n",
    "blob.download_to_filename(f'{directory}/operative_config.gin')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{directory}/checkpoint', 'w') as fopen:\n",
    "    fopen.write(f'model_checkpoint_path: \"model.ckpt-{model}\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from b2sdk.v1 import *\n",
    "info = InMemoryAccountInfo()\n",
    "b2_api = B2Api(info)\n",
    "application_key_id = b2_application_key_id\n",
    "application_key = b2_application_key\n",
    "b2_api.authorize_account(\"production\", application_key_id, application_key)\n",
    "file_info = {'how': 'good-file'}\n",
    "b2_bucket = b2_api.get_bucket_by_name('malaya-model')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7eff143a6e48>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tar = 't5-tiny-paraphrase-2021-07-31.tar.gz'\n",
    "os.system(f'tar -czvf {tar} {directory}')\n",
    "outPutname = f'finetuned/{tar}'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=tar,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "os.system(f'rm {tar}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import t5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = t5.models.MtfModel(\n",
    "    model_dir=directory,\n",
    "    tpu=None,\n",
    "    tpu_topology=None,\n",
    "    model_parallelism=1,\n",
    "    batch_size=1,\n",
    "    sequence_length={\"inputs\": 256, \"targets\": 256},\n",
    "    learning_rate_schedule=0.003,\n",
    "    save_checkpoints_steps=5000,\n",
    "    keep_checkpoint_max=3,\n",
    "    iterations_per_loop=100,\n",
    "    mesh_shape=\"model:1,batch:1\", \n",
    "    mesh_devices=[\"cpu:0\"]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf output/*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using config: {'_model_dir': 't5-tiny-paraphrase', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': None, '_session_config': allow_soft_placement: true\n",
      "isolate_session_state: true\n",
      ", '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': None, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_experimental_max_worker_delay_secs': None, '_session_creation_timeout_secs': 7200, '_service': None, '_cluster_spec': None, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1, '_tpu_config': TPUConfig(iterations_per_loop=100, num_shards=None, num_cores_per_replica=1, per_host_input_for_training=4, tpu_job_name=None, initial_infeed_sleep_secs=None, input_partition_dims=None, eval_training_input_configuration=2, experimental_host_call_every_n_steps=1), '_cluster': <tensorflow.python.distribute.cluster_resolver.tpu_cluster_resolver.TPUClusterResolver object at 0x7efe0043c7b8>}\n",
      "INFO:tensorflow:_TPUContext: eval_on_tpu True\n",
      "WARNING:tensorflow:eval_on_tpu ignored because use_tpu is False.\n"
     ]
    }
   ],
   "source": [
    "import gin\n",
    "\n",
    "from t5.data import sentencepiece_vocabulary\n",
    "\n",
    "DEFAULT_SPM_PATH = 'prepare/sp10m.cased.ms-en.model'\n",
    "DEFAULT_EXTRA_IDS = 100\n",
    "model_dir = directory\n",
    "\n",
    "def get_default_vocabulary():\n",
    "    return sentencepiece_vocabulary.SentencePieceVocabulary(\n",
    "      DEFAULT_SPM_PATH, DEFAULT_EXTRA_IDS)\n",
    "\n",
    "with gin.unlock_config():\n",
    "    gin.parse_config_file(t5.models.mtf_model._operative_config_path(model_dir))\n",
    "    gin.bind_parameter(\"Bitransformer.decode.beam_size\", 1)\n",
    "    gin.bind_parameter(\"Bitransformer.decode.temperature\", 0)\n",
    "    gin.bind_parameter(\"utils.get_variable_dtype.slice_dtype\", \"float32\")\n",
    "    gin.bind_parameter(\n",
    "        \"utils.get_variable_dtype.activation_dtype\", \"float32\")\n",
    "    \n",
    "vocabulary = t5.data.SentencePieceVocabulary(DEFAULT_SPM_PATH)\n",
    "estimator = model.estimator(vocabulary, disable_tpu=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1020200, 'model.ckpt-1020200', 't5-tiny-paraphrase/model.ckpt-1020200')"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "checkpoint_step = t5.models.mtf_model._get_latest_checkpoint_from_dir(model_dir)\n",
    "model_ckpt = \"model.ckpt-\" + str(checkpoint_step)\n",
    "checkpoint_path = os.path.join(model_dir, model_ckpt)\n",
    "checkpoint_step, model_ckpt, checkpoint_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1630: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Running infer on CPU\n",
      "INFO:tensorflow:feature inputs : Tensor(\"Reshape:0\", shape=(1, 256), dtype=int32)\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/mesh_tensorflow/transformer/utils.py:427: Print (from tensorflow.python.ops.logging_ops) is deprecated and will be removed after 2018-08-20.\n",
      "Instructions for updating:\n",
      "Use tf.print instead of tf.Print. Note that tf.print returns a no-output operator that directly prints the output. Outside of defuns or eager mode, this operator will not be executed unless it is directly specified in session.run or used as a control dependency for other operators. This is only a concern in graph mode. Below is an example of how to ensure tf.print executes in graph mode:\n",
      "\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "WARNING:tensorflow:Using default tf glorot_uniform_initializer for variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias  The initialzer will guess the input and output dimensions  based on dimension order.\n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/relative_attention_bias size 384          slice_size 384          Shape[heads=12, buckets=32]                                 \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/k                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/o                size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/q                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/EncDecAttention/v                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_000/layer_002/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/k                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/o                size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/q                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/EncDecAttention/v                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_001/layer_002/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/k                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/o                size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/q                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/EncDecAttention/v                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_002/layer_002/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/k                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/o                size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/q                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/EncDecAttention/v                size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable decoder/block_003/layer_002/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable decoder/final_layer_norm/scale                               size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/relative_attention_bias size 384          slice_size 384          Shape[heads=12, buckets=32]                                 \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_000/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_001/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_002/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/k                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/o                  size 294912       slice_size 294912       Shape[heads=768, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/q                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/SelfAttention/v                  size 294912       slice_size 294912       Shape[d_model=384, heads=768]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_000/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wi/kernel         size 516096       slice_size 516096       Shape[d_model=384, d_ff=1344]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/DenseReluDense/wo/kernel         size 516096       slice_size 516096       Shape[d_ff=1344, d_model=384]                               \n",
      "INFO:tensorflow:Variable encoder/block_003/layer_001/layer_norm/scale                 size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable encoder/final_layer_norm/scale                               size 384          slice_size 384          Shape[d_model=384]                                          \n",
      "INFO:tensorflow:Variable shared/embedding                                             size 12337152     slice_size 12337152     Shape[vocab=32128, d_model=384]                             \n",
      "INFO:tensorflow:Trainable Variables            count: 89      Total size: 34759680         Total slice_size: 34759680       \n",
      "INFO:tensorflow:All Variables                  count: 89      Total size: 34759680         Total slice_size: 34759680       \n",
      "INFO:tensorflow:Counters:\n",
      "einsum: 1.67e+10\n",
      "einsum_unique: 1.67e+10\n",
      "output: 2.13e+08\n",
      " output/AddOperation: 3.79e+07\n",
      " output/BinaryOpWithBroadcasting: 1.31e+06\n",
      " output/Constant: 1.57e+06\n",
      " output/EinsumOperation: 4.72e+07\n",
      " output/ImportOperation: 295\n",
      " output/MinMaxOperation: 7.87e+05\n",
      " output/OneHotOperation: 3.32e+07\n",
      " output/RangeOperation: 512\n",
      " output/ReduceOperation: 7.94e+04\n",
      " output/ReshapeOperation: 1.27e+07\n",
      " output/ScalarAddOperation: 1.05e+06\n",
      " output/ScalarMultiplyOperation: 2.46e+06\n",
      " output/ShiftOperation: 256\n",
      " output/SlicewiseOperation: 2.92e+07\n",
      " output/StopGradient: 9.44e+06\n",
      " output/Variable: 3.48e+07\n",
      " output/WhileLoopOperation: 1.57e+06\n",
      "output_unique: 2.13e+08\n",
      " output_unique/AddOperation: 3.79e+07\n",
      " output_unique/BinaryOpWithBroadcasting: 1.31e+06\n",
      " output_unique/Constant: 1.57e+06\n",
      " output_unique/EinsumOperation: 4.72e+07\n",
      " output_unique/ImportOperation: 295\n",
      " output_unique/MinMaxOperation: 7.87e+05\n",
      " output_unique/OneHotOperation: 3.32e+07\n",
      " output_unique/RangeOperation: 512\n",
      " output_unique/ReduceOperation: 7.94e+04\n",
      " output_unique/ReshapeOperation: 1.27e+07\n",
      " output_unique/ScalarAddOperation: 1.05e+06\n",
      " output_unique/ScalarMultiplyOperation: 2.46e+06\n",
      " output_unique/ShiftOperation: 256\n",
      " output_unique/SlicewiseOperation: 2.92e+07\n",
      " output_unique/StopGradient: 9.44e+06\n",
      " output_unique/Variable: 3.48e+07\n",
      " output_unique/WhileLoopOperation: 1.57e+06\n",
      "variables: 3.48e+07\n",
      " variables/trainable: 3.48e+07\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/ops/array_ops.py:1475: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "WARNING:tensorflow:From /home/husein/.local/lib/python3.6/site-packages/tensorflow_core/python/saved_model/signature_def_utils_impl.py:201: build_tensor_info (from tensorflow.python.saved_model.utils_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.utils.build_tensor_info or tf.compat.v1.saved_model.build_tensor_info.\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Classify: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Regress: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Predict: ['serving_default']\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Train: None\n",
      "INFO:tensorflow:Signatures INCLUDED in export for Eval: None\n",
      "INFO:tensorflow:Restoring parameters from t5-tiny-paraphrase/model.ckpt-1020200\n",
      "INFO:tensorflow:Assets added to graph.\n",
      "INFO:tensorflow:No assets to write.\n",
      "INFO:tensorflow:SavedModel written to: output/temp-b'1627712776'/saved_model.pb\n"
     ]
    }
   ],
   "source": [
    "from mesh_tensorflow.transformer import dataset as transformer_dataset\n",
    "\n",
    "def serving_input_fn():\n",
    "    inputs = tf.placeholder(\n",
    "            dtype=tf.string,\n",
    "            shape=[None],\n",
    "            name=\"inputs\")\n",
    "\n",
    "    batch_size = tf.shape(inputs)[0]\n",
    "    padded_inputs = tf.pad(inputs, [(0, tf.mod(-tf.size(inputs), batch_size))])\n",
    "    dataset = tf.data.Dataset.from_tensor_slices(padded_inputs)\n",
    "    dataset = dataset.map(lambda x: {\"inputs\": x})\n",
    "    dataset = transformer_dataset.encode_all_features(dataset, vocabulary)\n",
    "    dataset = transformer_dataset.pack_or_pad(\n",
    "        dataset=dataset,\n",
    "        length=model._sequence_length,\n",
    "        pack=False,\n",
    "        feature_keys=[\"inputs\"]\n",
    "    )\n",
    "    dataset = dataset.batch(tf.cast(batch_size, tf.int64))\n",
    "    features = tf.data.experimental.get_single_element(dataset)\n",
    "    return tf.estimator.export.ServingInputReceiver(\n",
    "        features=features, receiver_tensors=inputs)\n",
    "\n",
    "out = estimator.export_saved_model('output', serving_input_fn, checkpoint_path=checkpoint_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-15-5b89b6a20c22>:7: load (from tensorflow.python.saved_model.loader_impl) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "This function will only be available through the v1 compatibility library as tf.compat.v1.saved_model.loader.load or tf.compat.v1.saved_model.load. There will be a new function for importing SavedModels in Tensorflow 2.0.\n",
      "INFO:tensorflow:Restoring parameters from output/1627712776/variables/variables\n"
     ]
    }
   ],
   "source": [
    "config = tf.ConfigProto()\n",
    "config.allow_soft_placement = True\n",
    "sess = tf.Session(config = config)\n",
    "meta_graph_def = tf.saved_model.loader.load(\n",
    "        sess,\n",
    "        [tf.saved_model.tag_constants.SERVING],\n",
    "        out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'tiny-paraphrase/model.ckpt'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "saver = tf.train.Saver(tf.trainable_variables())\n",
    "saver.save(sess, 'tiny-paraphrase/model.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "strings = [\n",
    "    n.name\n",
    "    for n in tf.get_default_graph().as_graph_def().node\n",
    "    if ('encoder' in n.op\n",
    "    or 'decoder' in n.name\n",
    "    or 'shared' in n.name\n",
    "    or 'inputs' in n.name\n",
    "    or 'output' in n.name\n",
    "    or 'SentenceTokenizer' in n.name\n",
    "    or 'self/Softmax' in n.name)\n",
    "    and 'adam' not in n.name\n",
    "    and 'Assign' not in n.name\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze_graph(model_dir, output_node_names):\n",
    "\n",
    "    if not tf.gfile.Exists(model_dir):\n",
    "        raise AssertionError(\n",
    "            \"Export directory doesn't exists. Please specify an export \"\n",
    "            'directory: %s' % model_dir\n",
    "        )\n",
    "\n",
    "    checkpoint = tf.train.get_checkpoint_state(model_dir)\n",
    "    input_checkpoint = checkpoint.model_checkpoint_path\n",
    "\n",
    "    absolute_model_dir = '/'.join(input_checkpoint.split('/')[:-1])\n",
    "    output_graph = absolute_model_dir + '/frozen_model.pb'\n",
    "    clear_devices = True\n",
    "    with tf.Session(graph = tf.Graph()) as sess:\n",
    "        saver = tf.train.import_meta_graph(\n",
    "            input_checkpoint + '.meta', clear_devices = clear_devices\n",
    "        )\n",
    "        saver.restore(sess, input_checkpoint)\n",
    "        output_graph_def = tf.graph_util.convert_variables_to_constants(\n",
    "            sess,\n",
    "            tf.get_default_graph().as_graph_def(),\n",
    "            output_node_names,\n",
    "        )\n",
    "        with tf.gfile.GFile(output_graph, 'wb') as f:\n",
    "            f.write(output_graph_def.SerializeToString())\n",
    "        print('%d ops in the final graph.' % len(output_graph_def.node))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Restoring parameters from tiny-paraphrase/model.ckpt\n",
      "INFO:tensorflow:Froze 144 variables.\n",
      "INFO:tensorflow:Converted 144 variables to const ops.\n",
      "4840 ops in the final graph.\n"
     ]
    }
   ],
   "source": [
    "freeze_graph('tiny-paraphrase', strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import struct\n",
    "\n",
    "unknown = b'\\xff\\xff\\xff\\xff'\n",
    "\n",
    "def load_graph(frozen_graph_filename):\n",
    "    with tf.gfile.GFile(frozen_graph_filename, 'rb') as f:\n",
    "        graph_def = tf.GraphDef()\n",
    "        graph_def.ParseFromString(f.read())\n",
    "        \n",
    "    for node in graph_def.node:\n",
    "        \n",
    "        if node.op == 'RefSwitch':\n",
    "          node.op = 'Switch'\n",
    "          for index in xrange(len(node.input)):\n",
    "            if 'moving_' in node.input[index]:\n",
    "              node.input[index] = node.input[index] + '/read'\n",
    "        elif node.op == 'AssignSub':\n",
    "          node.op = 'Sub'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'AssignAdd':\n",
    "          node.op = 'Add'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "        elif node.op == 'Assign':\n",
    "          node.op = 'Identity'\n",
    "          if 'use_locking' in node.attr: del node.attr['use_locking']\n",
    "          if 'validate_shape' in node.attr: del node.attr['validate_shape']\n",
    "          if len(node.input) == 2:\n",
    "            node.input[0] = node.input[1]\n",
    "            del node.input[1]\n",
    "            \n",
    "        if 'Reshape/shape' in node.name or 'Reshape_1/shape' in node.name:\n",
    "            b = node.attr['value'].tensor.tensor_content\n",
    "            arr_int = [int.from_bytes(b[i:i + 4], 'little') for i in range(0, len(b), 4)]\n",
    "            if len(arr_int):\n",
    "                arr_byte = [unknown] + [struct.pack('<i', i) for i in arr_int[1:]]\n",
    "                arr_byte = b''.join(arr_byte)\n",
    "                node.attr['value'].tensor.tensor_content = arr_byte\n",
    "            \n",
    "            if len(node.attr['value'].tensor.int_val):\n",
    "                node.attr['value'].tensor.int_val[0] = -1\n",
    "    \n",
    "    with tf.Graph().as_default() as graph:\n",
    "        tf.import_graph_def(graph_def)\n",
    "    return graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "g = load_graph('tiny-paraphrase/frozen_model.pb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.Session(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.21 s, sys: 371 ms, total: 2.58 s\n",
      "Wall time: 621 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(1, 256)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "o_ = test_sess.run(o, feed_dict = {i: ['parafrasa: Speaker Dewan Rakyat Azhar Harun menegur ahli-ahli parlimen pembangkang selepas kira-kira sejam mereka menyatakan rasa tidak puas hati pada persidangan khas parlimen hari ini.',]})\n",
    "o_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sentencepiece as spm\n",
    "sp_model = spm.SentencePieceProcessor()\n",
    "sp_model.Load(DEFAULT_SPM_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 Speaker Dewan Rakyat Azhar Harun menegur MPs 's MPs Sue Harun rasa tidak sedap hati.\n"
     ]
    }
   ],
   "source": [
    "for k in range(len(o_)):\n",
    "    print(k, sp_model.DecodeIds(o_[k].tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.tools.graph_transforms import TransformGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "transforms = ['add_default_attributes',\n",
    "             'remove_nodes(op=Identity, op=CheckNumerics)',\n",
    "             'fold_batch_norms',\n",
    "             'fold_old_batch_norms',\n",
    "             'quantize_weights(minimum_size=1536000)',\n",
    "             #'quantize_weights(fallback_min=-10240, fallback_max=10240)',\n",
    "             'strip_unused_nodes',\n",
    "             'sort_by_execution_order']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-37-fe3ae91840e6>:3: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.gfile.GFile.\n"
     ]
    }
   ],
   "source": [
    "pb = 'tiny-paraphrase/frozen_model.pb'\n",
    "input_graph_def = tf.GraphDef()\n",
    "with tf.gfile.FastGFile(pb, 'rb') as f:\n",
    "    input_graph_def.ParseFromString(f.read())\n",
    "    \n",
    "transformed_graph_def = TransformGraph(input_graph_def, \n",
    "       ['inputs'],\n",
    "       ['SelectV2_3'], transforms)\n",
    "\n",
    "with tf.gfile.GFile(f'{pb}.quantized', 'wb') as f:\n",
    "    f.write(transformed_graph_def.SerializeToString())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(<tf.Tensor 'import/inputs:0' shape=(?,) dtype=string>,\n",
       " <tf.Tensor 'import/SelectV2_3:0' shape=(?, 256) dtype=int32>)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g = load_graph('tiny-paraphrase/frozen_model.pb.quantized')\n",
    "i = g.get_tensor_by_name('import/inputs:0')\n",
    "o = g.get_tensor_by_name('import/SelectV2_3:0')\n",
    "i, o"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_sess = tf.InteractiveSession(graph = g)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7efe5a8fb550>"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'tiny-paraphrase/frozen_model.pb.quantized'\n",
    "outPutname = 'paraphrase-v2/tiny-t5-quantized/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<b2sdk.file_version.FileVersionInfo at 0x7efd783e6d30>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "file = 'tiny-paraphrase/frozen_model.pb'\n",
    "outPutname = 'paraphrase-v2/tiny-t5/model.pb'\n",
    "b2_bucket.upload_local_file(\n",
    "    local_file=file,\n",
    "    file_name=outPutname,\n",
    "    file_infos=file_info,\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
